package com.sys.ai.service.abstracts;

import com.sys.ai.domain.AiEmbeddingDb;
import com.sys.ai.domain.AiKnowledge;
import com.sys.ai.domain.AiKnowledgeFragmentation;
import com.sys.ai.enums.VectorDBEnums;
import com.sys.ai.service.IAiEmbeddingDbService;
import com.sys.ai.service.IAiKnowledgeService;
import com.sys.ai.service.IDbOptionService;
import com.sys.ai.service.impl.DbOptionPostgreServiceImpl;
import com.sys.ai.vo.KnowledgeFragmentationVO;
import com.sys.common.utils.EnumUtil;
import com.sys.common.utils.spring.SpringUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * @author LGY
 * Create by 2025/7/1 10:03
 */
@Slf4j
public abstract class DbOptionAbstract implements IDbOptionService {

    protected static final String fragmentationField = "fragmentation";
    protected static final String fragmentationIdField = "fragmentation_id";
    protected static final String fragmentationEmbField = "fragmentation_emb";

    @Autowired
    protected IAiKnowledgeService aiKnowledgeService;
    @Autowired
    protected IAiEmbeddingDbService aiEmbeddingDbService;

    @Override
    public void createAiKnowledge(String schemeName, AiEmbeddingDb aiEmbeddingDb) {
        Objects.requireNonNull(getExecuteClass(aiEmbeddingDb)).createAiKnowledge(schemeName, aiEmbeddingDb);
    }

    @Override
    public void saveAiKnowledge(List<AiKnowledgeFragmentation> aiKnowledgeFragmentationList, String schemeName, AiEmbeddingDb aiEmbeddingDb) {
        Objects.requireNonNull(getExecuteClass(aiEmbeddingDb)).saveAiKnowledge(aiKnowledgeFragmentationList, schemeName, aiEmbeddingDb);
    }

    @Override
    public List<KnowledgeFragmentationVO> query(Float[] questionEmbedding, List<String> knowledgeIds) {
        List<KnowledgeFragmentationVO> data = new ArrayList<>();
        List<AiKnowledge> aiKnowledgeQaList = aiKnowledgeService.listByIds(knowledgeIds);
        aiKnowledgeQaList.forEach(item -> {
            try {
                AiEmbeddingDb aiEmbeddingDb = aiEmbeddingDbService.getById(item.getEmbeddingDbId());
                Assert.notNull(aiEmbeddingDb, "向量数据库不存在");
                List<KnowledgeFragmentationVO> res = Objects.requireNonNull(getExecuteClass(aiEmbeddingDb)).query(questionEmbedding, item,
                        aiEmbeddingDb);
                data.addAll(res);
            } catch (Exception e) {
                log.error("查询失败", e);
            }
        });
        if (CollectionUtils.isEmpty(data)) {
            KnowledgeFragmentationVO knowledgeFragmentationVO = new KnowledgeFragmentationVO();
            knowledgeFragmentationVO.setFragmentation("未从知识库中找到答案，请换个问题试试")
                    .setHasErrorTip(true);
            data.add(knowledgeFragmentationVO);
        }
        return data;
    }

    private IDbOptionService getExecuteClass(AiEmbeddingDb aiEmbeddingDb) {
        VectorDBEnums vectorDBEnums = EnumUtil.getEnumByCode(aiEmbeddingDb.getType(), VectorDBEnums.class);
        Assert.notNull(vectorDBEnums, "暂不支持的向量数据库类型");
        return switch (vectorDBEnums) {
            case POSTGRE -> SpringUtils.getBean(DbOptionPostgreServiceImpl.class);
            default -> null;
        };
    }

}
